#include <stdlib.h>
#include <stdio.h>
#include <stdint.h>
#include <string.h>
#include <ctype.h>
#include <time.h>


// L'état interne
struct {
  uint8_t i, j;
  uint8_t S[256];
} rc4;

inline void swap()
{
  uint8_t tmp = rc4.S[rc4.i];
  rc4.S[rc4.i] = rc4.S[rc4.j];
  rc4.S[rc4.j] = tmp;
}

// Initialisation de RC4 (Key Scheduling Algorithm)
void rc4_ksa(uint8_t *key, uint8_t len)
{
  rc4.i = rc4.j = 0;
  do {
    rc4.S[rc4.i] = rc4.i;
  } while (++rc4.i);
  do {
    rc4.j += rc4.S[rc4.i] + key[rc4.i % len];
    swap();
  } while (++rc4.i);
  rc4.j = 0;
}

// Génération d'un octet pseudo-aléatoire (Pseudo-Random Generation Algorithm)
uint8_t rc4_prga()
{
  rc4.j += rc4.S[++rc4.i];
  swap();
  return rc4.S[(rc4.S[rc4.i] + rc4.S[rc4.j]) & 0xff];
}



void usage()
{
  fprintf(stderr,
          "Usage : wepgen [-h] [-n NUM-MSG] [-s SEED] WEP-KEY\n"
          "Options et arguments :\n"
          "    -h           affiche uniquement ce message d'aide\n"
          "    -n NUM-MSG   génère NUM-MSG messages ; si cette option n'est "
          "pas spécifiée,\n"
          "                 génère des messages jusqu'à être interrompu par "
          "Ctrl-C\n"
          "    -s SEED      utilise SEED pour initialiser le générateur "
          "aléatoire\n"
          "    WEP-KEY      clé WEP de 40 ou 104 bits, en hexadécimal ;\n"
          "                 par exemple : d34db3ef42\n");
}

int main(int argc, char **argv)
{
  long int nmsg = -1, i;
  uint8_t  snap[3] = { 0xaa, 0xaa, 0x03 }, key[16], buf[6], len = 0, j;
  intmax_t seed = time(NULL);

  // On parse les options en ligne de commande
  for (--argc, ++argv; argc; --argc, ++argv) {
    if (!strcmp(*argv, "-h")) {
      usage();
      return EXIT_SUCCESS;
    }
    else if (!strcmp(*argv, "-n")) {
      char *ptr;
      --argc;
      nmsg = strtol(*++argv, &ptr, 10);
      if (*ptr || nmsg < 0) {
        fprintf(stderr, "Nombre de messages invalide : `%s'.\n", *argv);
        usage();
        return EXIT_FAILURE;
      }
    }
    else if (!strcmp(*argv, "-s")) {
      char *ptr;
      --argc;
      seed = strtol(*++argv, &ptr, 10);
      if (*ptr || seed < 0) {
        fprintf(stderr, "Graine du générateur aléatoire invalide : "
                        "`%s'.\n", *argv);
        usage();
        return EXIT_FAILURE;
      }
    }
    else if (**argv == '-') {
      fprintf(stderr, "Option non supportée : `%s'.\n", *argv);
      usage();
      return EXIT_FAILURE;
    }
    else if (len) {
      fprintf(stderr, "Clé WEP déjà définie.\n");
      usage();
      return EXIT_FAILURE;
    }
    else {
      uint8_t i;
      memset(key, 0x00, 16);
      for (i = 0, len = 3; i < 26 && argv[0][i]; len += i++ & 0x1) {
        uint8_t c = tolower(argv[0][i]);
        if (c < '0' || (c > '9' && c < 'a') || c > 'f')
          break;
        key[len] = (key[len] << 4) | (c >= 'a' ? 0xa + c - 'a' : c - '0');
      }
      if ((i != 10 && i != 26) || argv[0][i]) {
        fprintf(stderr, "Clé WEP invalide : `%s'.\n", *argv);
        usage();
        return EXIT_FAILURE;
      }
    }
  }

  if (!len) {
    fprintf(stderr, "Clé WEP non définie.\n");
    usage();
    return EXIT_FAILURE;
  }

  srand(seed);
  
  // Génération des IV, chiffrement de l'en-tête SNAP, et envoi des
  // paquets WEP sur STDOUT
  for (i = 0; nmsg < 0 || i < nmsg; ++i) {
    for (j = 0; j < 3; ++j)
      key[j] = rand();
    rc4_ksa(key, len);
    memcpy(buf, key, 3);
    for (j = 0; j < 3; ++j)
      buf[3+j] = snap[j] ^ rc4_prga();
    if (!fwrite(buf, 1, 6, stdout))
      break;
  }

  return EXIT_SUCCESS;
}
